#! /usr/bin/env python3
# -*- coding: utf-8 -*-

import uvicorn
from fastapi import FastAPI, Body
from fastapi.responses import JSONResponse
from typing import Dict
app = FastAPI()
from transformers import AutoTokenizer, AutoModel
model_name = "THUDM/chatglm3-6b"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name, trust_remote_code=True).half().cpu()
model = model.eval()

@app.post("/chat")
def f1(data: Dict):
    query = data["query"]
    history = data["history"]
    if history == "":
        history = []

    response, history = model.chat(tokenizer, query, history=history, top_p=0.95, temperature=0.95)
    response = {"response":response, "history":history}
    return JSONResponse(content=response) 

if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=7866)